package com.xxl.job.executor.service.jobhandler;

import com.xxl.job.core.biz.model.ReturnT;
import com.xxl.job.core.context.XxlJobContext;
import com.xxl.job.core.handler.annotation.XxlJob;
import com.xxl.job.executor.mybatis.service.JobUserService;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.batch.core.*;
import org.springframework.batch.core.configuration.annotation.EnableBatchProcessing;
import org.springframework.batch.core.configuration.annotation.JobBuilderFactory;
import org.springframework.batch.core.configuration.annotation.StepBuilderFactory;
import org.springframework.batch.core.job.builder.FlowBuilder;
import org.springframework.batch.core.job.flow.Flow;
import org.springframework.batch.core.launch.JobLauncher;
import org.springframework.batch.core.launch.support.RunIdIncrementer;
import org.springframework.batch.core.scope.context.ChunkContext;
import org.springframework.batch.core.step.tasklet.Tasklet;
import org.springframework.batch.repeat.RepeatStatus;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.SimpleAsyncTaskExecutor;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Date;

import static com.xxl.job.core.handler.IJobHandler.SUCCESS;

/**
 * @author: 吴多强
 * @create: 2020/7/23 16:54
 * @description: 使用Split实现任务并行执行
 **/
@Component
@Configuration
@EnableBatchProcessing
public class ShardingSpringBatchSplit {

    private static Logger logger = LoggerFactory.getLogger(ShardingSpringBatchSplit.class);
    @Resource
    private JobUserService jobUserService;

    @Autowired
    private JobLauncher jobLauncher;

    @Autowired
    private JobBuilderFactory jobBuilderFactory;

    @Autowired
    private StepBuilderFactory stepBuilderFactory;

    private int index, total, all, avg;

    @XxlJob("ShardingSpringBatchSplit")
    public ReturnT<String> shardingJobHandler(String param) throws Exception {
        // 当前分片序号
        index = XxlJobContext.getXxlJobContext().getShardIndex() + 1;
        // 总分片数
        total = XxlJobContext.getXxlJobContext().getShardTotal();
        // 总的分片任务
        all = Integer.parseInt(param);
        // 平均分片个数
        avg = all / total;
        logger.info("当前分片序号" + index + "总分片数" + total + "总的分片任务" + all + "平均分片个数" + avg);
        JobParameters jobParameters = new JobParametersBuilder().addDate("date", new Date())
                .toJobParameters();
        jobLauncher.run(splitJob(), jobParameters);
        return SUCCESS;
    }

    public Job splitJob() {
        return jobBuilderFactory.get("splitJobDemo")
                .incrementer(new RunIdIncrementer())
                .start(splitFlow1())
                // splitFlow1与splitFlow2并行执行
                .split(new SimpleAsyncTaskExecutor())
                .add(splitFlow2())
                .end()
                .build();
    }

    public Step splitStep1() {
        return stepBuilderFactory.get("splitStep1").tasklet(new Tasklet() {
            @Override
            public RepeatStatus execute(StepContribution stepContribution, ChunkContext chunkContext) throws Exception {
                logger.info("业务逻辑1" + "当前分片序号" + index + "总分片数" + total + "总的分片任务" + all + "平均分片个数" + avg);
                ShardingBatchSelect.selectJobUser(index, avg, jobUserService);
                return RepeatStatus.FINISHED;
            }
        }).build();
    }

    public Step splitStep2() {
        return stepBuilderFactory.get("splitStep2").tasklet(new Tasklet() {
            @Override
            public RepeatStatus execute(StepContribution stepContribution, ChunkContext chunkContext) throws Exception {
                logger.info("业务逻辑2" + "当前分片序号" + index + "总分片数" + total + "总的分片任务" + all + "平均分片个数" + avg);
                return RepeatStatus.FINISHED;
            }
        }).build();
    }

    public Step splitStep3() {
        return stepBuilderFactory.get("splitStep3").tasklet(new Tasklet() {
            @Override
            public RepeatStatus execute(StepContribution stepContribution, ChunkContext chunkContext) throws Exception {
                logger.info("业务逻辑3" + "当前分片序号" + index + "总分片数" + total + "总的分片任务" + all + "平均分片个数" + avg);
                return RepeatStatus.FINISHED;
            }
        }).build();
    }

    public Flow splitFlow1() {
        return new FlowBuilder<Flow>("splitFlow1")
                .start(splitStep1())
                .next(splitStep2())
                .build();
    }

    public Flow splitFlow2() {
        return new FlowBuilder<Flow>("splitFlow2")
                .start(splitStep3())
                .build();
    }
}
